import json
import os

import pandas as pd
import torch.cuda
from torch import optim
from torch.utils.data import DataLoader

from MeshDataset import *
from meshclsnet import *
import utils

if __name__ == '__main__':
    torch.cuda.set_device(utils.get_avail_gpu())
    torch.multiprocessing.set_start_method('spawn')

    with open('config.json', 'r') as f:
        config = json.load(f)

    mod = config["mod"]
    train_config = config["train"]

    train_list = os.path.join(config["list"]["list_path"] + mod, train_config["train_list"])
    val_list = os.path.join(config["list"]["list_path"] + mod, train_config["val_list"])

    model_path = train_config["model_path"] + mod
    model_name = train_config["model_name"] + mod
    checkpoint_name = train_config["checkpoint_name"]

    train_arg = config["train_arg"]

    num_classes = train_arg["num_classes"]
    num_channels = train_arg["num_channels"]
    num_epochs = train_arg["num_epochs"]
    num_workers = train_arg["num_workers"]
    train_batch_size = train_arg["train_batch_size"]
    val_batch_size = train_arg["val_batch_size"]
    num_batches_to_print = train_arg["num_batches_to_print"]

    if not os.path.exists(model_path):
        os.mkdir(model_path)

    training_dataset = Mesh_Dataset(data_list_path=train_list,
                                    num_seg_classes=15,
                                    patch_size=6000)
    val_dataset = Mesh_Dataset(data_list_path=train_list,
                               num_seg_classes=15,
                               patch_size=6000)
    train_loader = DataLoader(dataset=training_dataset,
                              batch_size=train_batch_size,
                              shuffle=True,
                              num_workers=num_workers)
    val_loader = DataLoader(dataset=val_dataset,
                            batch_size=val_batch_size,
                            shuffle=False,
                            num_workers=num_workers)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = MeshClsNet(num_classes=num_classes, num_channels=num_channels).to(device, dtype=torch.float)
    opt = optim.Adam(model.parameters(), amsgrad=True)

    train_mod = train_config["train_mod"]
    if train_mod == 1:
        checkpoint = torch.load(os.path.join(model_path, checkpoint_name), map_location='cpu')
        model.load_state_dict(checkpoint['model_state_dict'])
        opt.load_state_dict(checkpoint['optimizer_state_dict'])
        epoch_init = checkpoint['epoch']
        losses = checkpoint['losses']
        val_losses = checkpoint['val_losses']
        del checkpoint
        best_val_loss = 100.0
    else:
        losses = []
        val_losses = []
        epoch_init = 0
        best_val_loss = 100.0

    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True

    print('Training model...')
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        loss_epoch = 0.0
        for i_batch, batch_sample in enumerate(train_loader):
            inputs = batch_sample['cells'].to(device, dtype=torch.float)
            labels = batch_sample['labels'].to(device, dtype=torch.long)
            seg_matrix = batch_sample['seg_matrix'].to(device, dtype=torch.float)

            opt.zero_grad()

            outputs = model(inputs, seg_matrix)
            criteria = nn.CrossEntropyLoss()
            loss = criteria(outputs, labels)
            loss.backward()
            opt.step()

            running_loss += loss.item()
            loss_epoch += loss.item()

            if i_batch % num_batches_to_print == num_batches_to_print-1:
                print('[Epoch: {0}/{1}, Batch: {2}/{3}] loss: {4}'.format(epoch+1, num_epochs, i_batch+1, len(train_loader), running_loss/num_batches_to_print))
                running_loss = 0.0

        losses.append(loss_epoch/len(train_loader))
        loss_epoch = 0.0

        model.eval()
        with torch.no_grad():
            running_val_loss = 0.0
            val_loss_epoch = 0.0
            for i_batch, batch_val_sample in enumerate(val_loader):
                inputs = batch_val_sample['cells'].to(device, dtype=torch.float)
                labels = batch_val_sample['labels'].to(device, dtype=torch.long)
                seg_matrix = batch_val_sample['seg_matrix'].to(device, dtype=torch.float)

                outputs = model(inputs, seg_matrix)
                criteria = nn.CrossEntropyLoss()
                loss = criteria(outputs, labels)

                running_val_loss += loss.item()
                val_loss_epoch += loss.item()

                if i_batch % num_batches_to_print == num_batches_to_print-1:
                    print('[Epoch: {0}/{1}, Val batch: {2}/{3}] val_loss: {4}'.format(epoch+1, num_epochs, i_batch+1, len(val_loader), running_val_loss/num_batches_to_print))
                    running_val_loss = 0.0

            val_losses.append(val_loss_epoch/len(val_loader))
            val_loss_epoch = 0.0

            print('*****\nEpoch: {}/{}, loss{}, val_loss{}\n*****'.format(epoch+1, num_epochs, losses[-1], val_losses[-1]))

        torch.save({
            'epoch': epoch+1,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': opt.state_dict(),
            'losses': losses,
            'val_losses': val_losses
        }, model_path+checkpoint_name)

        if best_val_loss > val_losses[-1]:
            torch.save({
                'epoch': epoch + 1,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': opt.state_dict(),
                'losses': losses,
                'val_losses': val_losses
            }, model_path+'{}_best.tar'.format(model_name))

        pd_dict = {'loss': losses, 'val_loss': val_losses}
        stat = pd.DataFrame(pd_dict)
        stat.to_csv(os.path.join(model_path, 'losses_metrics_vs_epoch_' + mod + '.csv'))